import numpy as np
import tensorflow as tf
from data_helper import MedicineData
from models import ConvolutionNet, SemiConvolutionNet


def main():
    model = ConvolutionNet((128, 128, 3), drop_out=0.5)
    data = MedicineData()

    sess = tf.InteractiveSession()
    sess.run(tf.global_variables_initializer())
    '0.9054'
    epochs = 100
    for epoch in range(epochs):
        loss_lst = []
        for batch_x, batch_y in data.batch_iter(data_type='train', batch_size=50):
            train_dict = {model.inputs: batch_x, model.labels: batch_y}
            _, loss = sess.run([model.optimizer, model.loss], feed_dict=train_dict)
            loss_lst.append(loss)
        losses = np.mean(loss_lst)
        print('Epoch:{}/{}'.format(epoch, epochs),
              'Train losses: {:.4f}'.format(losses))
        if epoch % 5 == 0:
            acc_lst = []
            for test_x, test_y in data.batch_iter(data_type='test', batch_size=64):
                test_dict = {model.inputs: test_x, model.labels: test_y}
                acc = sess.run([model.accuracy], feed_dict=test_dict)
                acc_lst.append(acc[0])
            acc = np.mean(acc_lst)
            print('Epoch:{}/{}'.format(epoch, epochs),
                  'Test acc: {:.4f}'.format(acc))
    sess.close()


def train_semi():
    model = SemiConvolutionNet((128, 128, 3), drop_out=0.5)
    data = MedicineData()

    sess = tf.InteractiveSession()
    sess.run(tf.global_variables_initializer())
    # 先train none label
    auto_epochs = 100
    for auto_epoch in range(auto_epochs):
        auto_loss_lst = []
        for auto_batch_x, auto_batch_y in data.batch_iter(data_type='none', batch_size=50):
            auto_dict = {model.inputs: auto_batch_x, model.labels: auto_batch_y}
            _, auto_loss = sess.run([model.auto_optimizer, model.auto_loss], feed_dict=auto_dict)
            print(auto_loss)
            auto_loss_lst.append(auto_loss)
        losses = np.mean(auto_loss_lst)
        print('Epoch:{}/{}'.format(auto_epoch, auto_epochs),
              'Train losses: {:.4f}'.format(losses))
    # 正常train

    epochs = 300
    for epoch in range(epochs):
        loss_lst = []
        for batch_x, batch_y in data.batch_iter(data_type='train', batch_size=50):
            train_dict = {model.inputs: batch_x, model.labels: batch_y}
            _, loss = sess.run([model.optimizer, model.loss], feed_dict=train_dict)
            loss_lst.append(loss)
        losses = np.mean(loss_lst)
        print('Epoch:{}/{}'.format(epoch, epochs),
              'Train losses: {:.4f}'.format(losses))
        if epoch % 5 == 0:
            acc_lst = []
            for test_x, test_y in data.batch_iter(data_type='test', batch_size=64):
                test_dict = {model.inputs: test_x, model.labels: test_y}
                acc = sess.run([model.accuracy], feed_dict=test_dict)
                acc_lst.append(acc[0])
            acc = np.mean(acc_lst)
            print('Epoch:{}/{}'.format(epoch, epochs),
                  'Test acc: {:.4f}'.format(acc))
    sess.close()


if __name__ == '__main__':
    # main()
    train_semi()
